from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from dataset import TaoBaoDataset
from models.MyModel import *
from models.MyModel_cos import *
from models.MyModleConfig import *
import wandb
from config import *
import torch.optim as optim
import os
from sklearn.metrics import roc_auc_score
from numpy import mean
import warnings
from torch.utils.tensorboard import SummaryWriter
import seaborn as sns
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

warnings.filterwarnings('ignore')


def prepare_model_data(model_name, chose_epoch):
    test_target_item_file = '../data/taobao/final_test_data.csv'
    user_behavior_file = f'../data/taobao/active_user_actions.csv'
    n_users = None
    test_data = TaoBaoDataset(test_target_item_file, user_behavior_file, mode='test', n_users=n_users, q_num=q_num)
    test_loader = DataLoader(test_data, batch_size=6, shuffle=False)

    config = MyModelConfig(
        test_data.vocab_size_dic,
        HEADS,
        q_num=q_num,
        short_time=short_time,
        split_range=split_range,
        total_behavior_num=total_behavior_num,
        device_id=device_id,
        compressed=compressed,
        layers1=layers1,
        layers2=layers2,
        diff1=diff1,
        diff2=diff2
    )
    device = config.device

    if model_name =='myModel':
        model = MyModel(config)
    elif model_name == 'myModel_cos':
        model = MyModel_Cos(config)
    model = model.to(device)

    checkpoint = torch.load(f"checkpoints/{model_name}_{chose_epoch}.pth", weights_only=False)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    print(f"checkpoints/{model_name}_{chose_epoch}.pth模型数据准备完毕")
    return model, test_loader, config


def print_us_target_dot(model_name, model, loader, config):
    model.eval()
    print("------点积热力图------")
    device = config.device
    with torch.no_grad():
        for target_item, behaviors, global_interest_tokens,target in loader:
            target_item = target_item.to(device)
            behaviors = behaviors.to(device)
            target = target.to(device)
            global_interest_tokens = global_interest_tokens.to(device)
            pred, atte1, atte2, us, target_embed, ts = model(behaviors, global_interest_tokens, target_item)
            us = us.float().to(device)
            target_embed = target_embed.float().to(device)
            b = us.shape[0]
            y = torch.einsum("bqd,bad->bqa", [us, target_embed]).view(b, -1)
            if not os.path.exists("imgs"):
                os.mkdir("imgs/")
            plt.figure(figsize=(10, 100))
            sns.heatmap(y.cpu(), cmap='viridis', annot=True, fmt=".4f", cbar=True)
            plt.title("Heatmap using Seaborn")
            plt.xlabel("interest")
            plt.ylabel("target item")
            plt.savefig(
                f"imgs/{model_name}_{layers1}_{layers2}_{q_num}_{loss_list}_{lambda_list}_us_target_dot.png")
            break


def print_atte_score(model_name, model, loader, config):
    model.eval()
    print("------注意力得分热力图------")
    device = config.device
    with torch.no_grad():
        for target_item, behaviors, global_interest_tokens,target in loader:
            target_item = target_item.to(device)
            behaviors = behaviors.to(device)
            target = target.to(device)
            global_interest_tokens = global_interest_tokens.to(device)
            pred, atte1, atte2, us, target_embed, ts = model(behaviors, global_interest_tokens, target_item)
            us = us.float().to(device)
            atte1 = atte1.float().to(device)
            atte1 = atte1.mean(dim=1)
            atte2 = atte2.float().to(device)
            atte2 = atte2.mean(dim=1)
            b, q, e = atte1.shape
            # y = y[:10,:,:].reshape(10*q,e)
            atte1 = atte1.reshape(b * q, e)
            b, q, e = atte2.shape
            atte2 = atte2.reshape(b * q, e)

            if not os.path.exists("imgs"):
                os.mkdir("imgs/")
            plt.figure(figsize=(200, 100))
            sns.heatmap(atte1.cpu(), cmap='viridis', annot=True, fmt=".4f", cbar=True)
            plt.title("")
            plt.xlabel("short time item")
            plt.ylabel("target item")
            plt.savefig(
                f"imgs/{model_name}_{layers1}_{layers2}_{q_num}_{loss_list}_{lambda_list}_item_atte_score.png")
            

            plt.figure(figsize=(200, 100))
            sns.heatmap(atte2.cpu(), cmap='viridis', annot=True, fmt=".4f", cbar=True)
            plt.title("")
            plt.xlabel("interest")
            plt.ylabel("target item")
            plt.savefig(
                f"imgs/{model_name}_{layers1}_{layers2}_{q_num}_{loss_list}_{lambda_list}_interests_atte_score.png")
            break

def t_SNE_visualization(model_name, model, loader, config):
    model.eval()
    print("------embedding图------")
    device = config.device
    tsne = TSNE(n_components=2, random_state=42, perplexity=30, learning_rate=200)
    with torch.no_grad():
        for idx,(target_item, behaviors, global_interest_tokens,target) in enumerate(loader):
            target_item = target_item.to(device)
            behaviors = behaviors.to(device)
            target = target.to(device)
            global_interest_tokens = global_interest_tokens.to(device)
            pred, atte1, atte2, us, target_embed, ts = model(behaviors, global_interest_tokens, target_item)
            total = torch.concat([us[0], target_embed[:2].squeeze(1), ts[0]], dim = 0)
            total = tsne.fit_transform( total.cpu().numpy() )
            us = total[:model.q_num]
            target_embed = total[model.q_num: model.q_num+2]
            pos_target_embed = target_embed[0]
            neg_target_embed = target_embed[1]
            ts = total[model.q_num+2:]
            early_behaovir_embed = ts[-(model.behavior_group_num[2] // model.early_group_num):,:]
            middle_behavoir_embed = ts[ -(model.behavior_group_num[1] // model.middle_group_num + model.behavior_group_num[2] // model.early_group_num):-(model.behavior_group_num[2] // model.early_group_num), :]
            latest_behavoir_embed = ts[model.short_time:model.behavior_group_num[0], :]
            short_time_embed = ts[:model.short_time,:]
            if not os.path.exists("imgs"):
                os.mkdir("imgs/")
            plt.figure(figsize=(10, 10))
            plt.scatter(us[0,0], us[0,1], c= 'red',marker="^",s=200,label="interest1")
            plt.scatter(us[1,0], us[1,1], c= 'yellow',marker="^",s=200,label="interest2")
            plt.scatter(us[2,0], us[2,1], c= 'pink',marker="^",s=200,label="interest3")
            plt.scatter(us[3,0], us[3,1], c= 'orange',marker="^",s=200,label="interest4")
            plt.scatter(pos_target_embed[0], pos_target_embed[1], c='green', marker='o',s=200,label="pos_item")
            plt.scatter(neg_target_embed[0], neg_target_embed[1], c='gray', marker='o',s=200,label="neg_item")
            plt.scatter(early_behaovir_embed[:,0], early_behaovir_embed[:,1], c='blue', marker='s',alpha=0.2,label="early_behaviors")
            plt.scatter(middle_behavoir_embed[:,0], middle_behavoir_embed[:,1], c='blue', marker='s',alpha=0.5,label="middle_behaviors")
            plt.scatter(latest_behavoir_embed[:,0], latest_behavoir_embed[:,1], c='blue', marker='s',alpha=1,label="latest_behaviors")
            plt.scatter(short_time_embed[:,0], short_time_embed[:,1], c='purple', marker='s',alpha=1,label="latest_behaviors")
            plt.title(f"{model_name}_batch_{idx}_first_user_embedding")
            plt.xlabel("x")
            plt.ylabel("y")
            plt.legend()
            plt.savefig(
                f"imgs/{model_name}_{layers1}_{layers2}_{q_num}_{loss_list}_{lambda_list}_embedding_tsne_{idx}.png")
            if idx == 10:
                break


if __name__ == '__main__':
    model, test_loader, config = prepare_model_data(model_name, chose_epoch=9)
    print_us_target_dot(model_name, model, test_loader, config)
    print_atte_score(model_name,model,test_loader,config)
    t_SNE_visualization(model_name,model,test_loader,config)
